package com.ruoyi.kindergarten.websocket;

import org.springframework.stereotype.Component;

import javax.websocket.*;
import javax.websocket.server.ServerEndpoint;
import java.io.IOException;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;

@ServerEndpoint(value = "/ws/Iot")
@Component
public class IotWebSocketService {
    private static final ConcurrentHashMap<String, Session> SESSIONS  = new ConcurrentHashMap<String, Session>();
    private static final ConcurrentHashMap<String, Long> TIME_STAMP = new ConcurrentHashMap<>();
    private static final ConcurrentHashMap<String, AtomicInteger> MSG_COUNTER = new ConcurrentHashMap<>();
    private static final long TIME_WINDOW = 1000; // 1秒
    private static final int MAX_MSG_COUNT = 10;  // 1秒内最多10条消息

    @OnMessage
    public void onMessage(String message, Session session) {
        System.out.println(message);
        if (!allowSend(session.getId())){
            return;
        }
        broadcast(message);
    }
    @OnOpen
    public void onOpen(Session session) {
        long now = System.currentTimeMillis();
        MSG_COUNTER.put(session.getId(), new AtomicInteger(0));
        TIME_STAMP.put(session.getId(), now);
        SESSIONS.put(session.getId(), session);
    }
    @OnClose
    public void onClose(Session session) {
        SESSIONS.remove(session.getId());
        MSG_COUNTER.remove(session.getId());
        TIME_STAMP.remove(session.getId());
    }
    @OnError
    public void onError(Session session, Throwable throwable) {
        SESSIONS.remove(session.getId());
        MSG_COUNTER.remove(session.getId());
        TIME_STAMP.remove(session.getId());
        throwable.printStackTrace();
        System.out.println(this.getClass().getName() + " onError");
    }
    public void broadcast(String message) {
        for (Session value : SESSIONS.values()) {
            try{
                value.getBasicRemote().sendText(message);
            }catch (IOException e){
                e.printStackTrace();
            }
        }
    }

    public boolean allowSend(String sessionId){
        long now = System.currentTimeMillis();


        AtomicInteger counter  = MSG_COUNTER.get(sessionId);
        Long startTime = TIME_STAMP.get(sessionId);

        if(now - startTime > TIME_WINDOW){
            counter.set(0);
            TIME_STAMP.put(sessionId, now);
        }

        int currentCount = counter.incrementAndGet();
        return currentCount <= MAX_MSG_COUNT;
    }
}
